import cv2
import torch
import numpy as np
from PIL import Image, ImageOps, ImageSequence


def convert_non_white_to_black(mask_tensor):
    """
    将张量中除了纯白色(RGB = (1.0, 1.0, 1.0))以外的所有颜色转换为黑色(RGB = (0.0, 0.0, 0.0))。
    支持单通道和多通道张量。

    :param mask_tensor: 浮点型张量掩码，形状可以为(H, W)或(H, W, C)，其中C为颜色通道数。
    """
    # 确保输入是Tensor类型
    assert isinstance(mask_tensor, torch.Tensor), "Input should be a PyTorch Tensor."

    # 根据张量维度处理
    if mask_tensor.dim() == 2:  # 处理灰度图或单通道情况
        white_pixels = mask_tensor == 1.0
        result_tensor = torch.where(
            white_pixels, mask_tensor, torch.zeros_like(mask_tensor)
        )
    elif mask_tensor.dim() == 3:  # 多通道情况，假设最后一个维度为颜色通道
        white_pixels = (mask_tensor == 1.0).all(dim=-1)
        result_tensor = torch.where(
            white_pixels.unsqueeze(-1), mask_tensor, torch.zeros_like(mask_tensor)
        )
    else:
        raise ValueError("Unsupported tensor dimension. Expected 2D or 3D tensor.")

    return result_tensor


def tensor_to_image(tensor):
    """
    将二维的PyTorch Tensor转换为PIL图像。

    参数:
    - tensor: 一个形状为(H, W)或(C, H, W)的二维或三维Tensor，其中C是通道数（对于灰度图C=1，彩色图C=3）。

    返回:
    - 一个PIL.Image对象。
    """
    if tensor.max() <= 1 and tensor.min() >= 0:  # 值在[0,1]之间
        tensor = tensor * 255  # 转换到[0,255]区间
    tensor = tensor.clamp(0, 255).to(torch.uint8)  # 确保值在[0,255]并转换为无符号整型
    image_array = tensor.numpy()
    if image_array.ndim == 2:  # 确保是灰度图
        image = Image.fromarray(image_array, "L")  # 'L' 表示灰度图
    else:
        image = Image.fromarray(image_array)
    return image


def pil_to_tensor(image_array):
    """
    将numpy数组转换为PyTorch Tensor.

    参数:
    - image_array: numpy数组，代表图像数据。
    返回:
    - 转换后的 Tensor，形状为 (通道数, 高度, 宽度)，值在 [0.0, 1.0] 之间（浮点数）。
    """
    # 确保已经是浮点型且值在[0,1]之间，如果是其他范围需要归一化
    if image_array.dtype != np.float32:
        image_array = image_array.astype(np.float32) / 255.0

    # 根据图像类型处理
    if len(image_array.shape) == 2:  # 灰假设灰度图
        tensor = torch.from_numpy(image_array[None,])  # 添加通道维度
    elif len(image_array.shape) == 3:  # 假设RGB图
        tensor = torch.from_numpy(image_array.transpose((2, 0, 1)))  # CHW格式
    else:
        raise ValueError("Unsupported image shape.")
    return tensor


def add_soft_outline_and_white_background(
    pil_image,
    color_threshold=(255, 255, 255),
    line_thickness=1,
    line_color=(255, 255, 255),
    blur_kernel_size=(3, 3),
):
    """
    在指定颜色区域的边缘添加柔和的轮廓线，并将非轮廓区域的颜色变为黑色。

    :param pil_image: PIL灰度图片对象
    :param color_threshold: 需要识别的颜色阈值，默认为白色(255, 255, 255)
    :param line_thickness: 轮廓线的粗细
    :param line_color: 轮廓线的颜色
    :param blur_kernel_size: 高斯模糊的核大小，决定了轮廓柔和的程度
    """
    # 将PIL图片转换为OpenCV格式
    img = np.array(tensor_to_image(pil_image))

    if img is None:
        print("图片读取失败，请检查图片路径。")
        return

    # 创建掩码
    lower_bound = np.array(color_threshold, dtype=np.uint8)
    upper_bound = np.array(color_threshold, dtype=np.uint8)

    # 确保lower_bound和upper_bound的形状与img匹配
    if len(img.shape) == 2:  # 灰度图像
        lower_bound = np.array([color_threshold[0]], dtype=np.uint8)
        upper_bound = np.array([color_threshold[0]], dtype=np.uint8)
    elif len(img.shape) == 3 and img.shape[2] == 3:  # 彩色图像
        lower_bound = np.array(
            [color_threshold[0], color_threshold[1], color_threshold[2]], dtype=np.uint8
        )
        upper_bound = np.array(
            [color_threshold[0], color_threshold[1], color_threshold[2]], dtype=np.uint8
        )

    mask = cv2.inRange(img, lower_bound, upper_bound)

    # 膨胀操作
    kernel = np.ones((3, 3), np.uint8)
    mask_dilated = cv2.dilate(mask, kernel, iterations=1)

    # 找到轮廓
    contours, _ = cv2.findContours(
        mask_dilated - mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    # 准备绘制轮廓的图像副本
    img_with_outline = np.zeros_like(img)  # 创建全黑背景图像

    for contour in contours:
        # 绘制轮廓线
        cv2.drawContours(img_with_outline, [contour], 0, line_color, line_thickness)

    img_with_outline_tensor = torch.from_numpy(img_with_outline).float()
    return img_with_outline_tensor
